#!/usr/bin/env python3

import sys
import time
import json

try:
    from p4utils.utils.helper import load_topo
    from p4utils.utils.sswitch_thrift_API import SimpleSwitchThriftAPI
except ImportError:
    print("警告: 无法导入P4Utils，尝试使用备用方法...")

class SourceRoutingController:
    def __init__(self, switch_name):
        self.switch_name = switch_name
        try:
            self.controller = SimpleSwitchThriftAPI(switch_name)
            self.topo = load_topo()
        except:
            print(f"无法连接到交换机 {switch_name}，使用模拟模式")
            self.controller = None

    def set_source_routing_entries(self):
        """设置源路由表项"""
        print(f"在 {self.switch_name} 上设置源路由表项...")

        if not self.controller:
            print(f"模拟模式: {self.switch_name} 源路由表项设置")
            return

        try:
            # 清除现有表项
            self.controller.table_clear("src_route_table")

            if self.switch_name == "s1":
                self.controller.table_add("src_route_table", "src_route_forward", ["1"], ["2"])
                print(f"  s1: 源路由 -> 端口2 (s11)")

            elif self.switch_name == "s11":
                self.controller.table_add("src_route_table", "src_route_forward", ["1"], ["2"])
                print(f"  s11: 源路由 -> 端口2 (s12)")

            elif self.switch_name == "s12":
                self.controller.table_add("src_route_table", "src_route_forward", ["1"], ["2"])
                print(f"  s12: 源路由 -> 端口2 (s2)")

            elif self.switch_name == "s21":
                self.controller.table_add("src_route_table", "src_route_forward", ["1"], ["2"])
                print(f"  s21: 源路由 -> 端口2 (s22)")

            elif self.switch_name == "s22":
                self.controller.table_add("src_route_table", "src_route_forward", ["1"], ["2"])
                print(f"  s22: 源路由 -> 端口2 (s2)")

            elif self.switch_name == "s2":
                self.controller.table_add("src_route_table", "src_route_forward", ["1"], ["1"])
                print(f"  s2: 源路由 -> 端口1 (h2)")

            print(f"✓ {self.switch_name} 源路由表项设置完成")

        except Exception as e:
            print(f"✗ {self.switch_name} 源路由表项设置失败: {e}")

    def set_ipv4_forwarding(self):
        """设置IPv4转发表项"""
        print(f"在 {self.switch_name} 上设置IPv4转发表项...")

        if not self.controller:
            print(f"模拟模式: {self.switch_name} IPv4表项设置")
            return

        try:
            # 清除现有表项
            self.controller.table_clear("ipv4_table")

            if self.switch_name == "s1":
                self.controller.table_add("ipv4_table", "ipv4_forward", ["10.0.0.2/32"], ["2"])
                self.controller.table_add("ipv4_table", "ipv4_forward", ["10.0.0.1/32"], ["1"])
                print(f"  s1: 10.0.0.2 -> 端口2, 10.0.0.1 -> 端口1")

            elif self.switch_name == "s2":
                self.controller.table_add("ipv4_table", "ipv4_forward", ["10.0.0.1/32"], ["2"])
                self.controller.table_add("ipv4_table", "ipv4_forward", ["10.0.0.2/32"], ["1"])
                print(f"  s2: 10.0.0.1 -> 端口2, 10.0.0.2 -> 端口1")

            elif self.switch_name == "s11":
                self.controller.table_add("ipv4_table", "ipv4_forward", ["10.0.0.2/32"], ["2"])
                print(f"  s11: 10.0.0.2 -> 端口2")

            elif self.switch_name == "s12":
                self.controller.table_add("ipv4_table", "ipv4_forward", ["10.0.0.2/32"], ["2"])
                print(f"  s12: 10.0.0.2 -> 端口2")

            elif self.switch_name == "s21":
                self.controller.table_add("ipv4_table", "ipv4_forward", ["10.0.0.2/32"], ["2"])
                print(f"  s21: 10.0.0.2 -> 端口2")

            elif self.switch_name == "s22":
                self.controller.table_add("ipv4_table", "ipv4_forward", ["10.0.0.2/32"], ["2"])
                print(f"  s22: 10.0.0.2 -> 端口2")

            print(f"✓ {self.switch_name} IPv4转发表项设置完成")

        except Exception as e:
            print(f"✗ {self.switch_name} IPv4转发表项设置失败: {e}")

def main():
    """主控制函数"""
    switches = ["s1", "s2", "s11", "s12", "s21", "s22"]

    print("=== 开始配置源路由和IPv4转发 ===")

    for switch_name in switches:
        try:
            print(f"\n--- 配置 {switch_name} ---")
            controller = SourceRoutingController(switch_name)
            controller.set_source_routing_entries()
            controller.set_ipv4_forwarding()
            time.sleep(0.5)
        except Exception as e:
            print(f"配置 {switch_name} 时出错: {e}")
            continue

    print("\n=== 所有交换机配置完成 ===")
    print("现在可以运行测试脚本: python source_routing_test.py")

if __name__ == "__main__":
    main()
